{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "mraa-body.ipynb",
      "private_outputs": true,
      "provenance": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/eyaler/avatars4all/blob/master/mraa-body.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cdO_RxQZLahB"
      },
      "source": [
        "# Demo for paper \"Motion Representations for Articulated Animation\"\n",
        "\n",
        "## **Full body Tai chi and TED models**\n",
        "\n",
        "### Made just a little bit more accessible by Eyal Gruss (https://eyalgruss.com, eyalgruss@gmail.com)\n",
        "\n",
        "##### Original project: https://snap-research.github.io/articulated-animation\n",
        "\n",
        "##### Original notebook: https://colab.research.google.com/github/AliaksandrSiarohin/articulated-animation/blob/main/demo.ipynb\n",
        "\n",
        "##### My avatars4all repository: https://github.com/eyaler/avatars4all\n",
        "##### A list of more generative tools: https://j.mp/generativetools"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UCMFMJV7K-ag",
        "cellView": "form"
      },
      "source": [
        "#@title Setup\n",
        " \n",
        "%cd /content\n",
        "#!curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash\n",
        "#!sudo apt-get install git-lfs\n",
        "#!git lfs install\n",
        "!pip install -U gdown\n",
        "!git clone https://github.com/eyaler/articulated-animation\n",
        "!gdown --id 1zt1cbLGiN6pVZi1Wn70lF29l1c2x6--V\n",
        "!gdown --id 1UXMd2x4G3uFLa-XPA5iGywWlcidXR6Yg\n",
        "!gdown --id 1y_eiItgdKr3Z0Zb7lCsQoCM2GyxAcya_\n",
        "!gdown --id 1X-BwTTwxTteesAxymrutxOy2qFW_jlWA\n",
        "!gdown --id 1gyU9H7ArRAK9cUp_4rGF582C35B4pNcE\n",
        "!mv /content/*.pth articulated-animation/checkpoints\n",
        "\n",
        "!pip install -U git+https://github.com/ytdl-org/youtube-dl\n",
        "!pip install imageio==2.9.0\n",
        "!pip install imageio-ffmpeg==0.4.5\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n5L7MTJcJxP2",
        "cellView": "form"
      },
      "source": [
        "#@title Choose model\n",
        "model = 'ted384' #@param ['taichi256', 'ted384', 'ted-youtube384']\n",
        "\n",
        "res = int(model[-3:])\n",
        "%cd /content/articulated-animation\n",
        "from demo import load_checkpoints\n",
        "generator, region_predictor, avd_network = load_checkpoints(config_path='config/%s.yaml'%model,\n",
        "                                                            checkpoint_path='checkpoints/%s.pth'%model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2iGy_8v-4BrO",
        "cellView": "form"
      },
      "source": [
        "#@title Get the Driver video and Avatar image from the web\n",
        "#@markdown 1. You can change the URLs to your **own** stuff!\n",
        "#@markdown 2. Alternatively, you can upload **local** files in the next cells\n",
        " \n",
        "video_url = 'https://www.youtube.com/watch?v=ZKHpkVXZFZ4' #@param {type:\"string\"}\n",
        "image_url = 'https://s3-eu-west-1.amazonaws.com/uploads.playbaamboozle.com/uploads/images/94473/1613432331_112044.jpeg' #@param {type:\"string\"}\n",
        " \n",
        "if video_url:\n",
        "  !rm -f /content/video.mp4\n",
        "  !youtube-dl --no-check-extensions --no-playlist -f 'bestvideo[ext=mp4][vcodec!*=av01][height<=360]+bestaudio[ext=m4a]/mp4[height<=360][vcodec!*=av01]/mp4[vcodec!*=av01]/mp4' \"$video_url\" --merge-output-format mp4 -o /content/video\n",
        "  !mv /content/video.mp4 /content/video \n",
        " \n",
        "if image_url:\n",
        "  !wget \"$image_url\" -O /content/image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y1xCcE9XoSy4",
        "cellView": "form"
      },
      "source": [
        "#@title Optionally upload local Driver video { run: \"auto\" }\n",
        "manually_upload_video = False #@param {type:\"boolean\"}\n",
        "if manually_upload_video:\n",
        "  from google.colab import files\n",
        "  import shutil\n",
        "\n",
        "  %cd /content/sample_data\n",
        "  try:\n",
        "    uploaded = files.upload()\n",
        "  except Exception as e:\n",
        "    %cd /content\n",
        "    raise e\n",
        "\n",
        "  for fn in uploaded:\n",
        "    shutil.move('/content/sample_data/'+fn, '/content/video')\n",
        "    break\n",
        "  %cd /content"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m-zK9rlDoS7l",
        "cellView": "form"
      },
      "source": [
        "#@title Optionally upload local Avatar image { run: \"auto\" }\n",
        "manually_upload_image = False #@param {type:\"boolean\"}\n",
        "if manually_upload_image:\n",
        "  from google.colab import files\n",
        "  import shutil\n",
        " \n",
        "  %cd /content/sample_data\n",
        "  try:\n",
        "    uploaded = files.upload()\n",
        "  except Exception as e:\n",
        "    %cd /content\n",
        "    raise e\n",
        "\n",
        "  for fn in uploaded:\n",
        "    shutil.move('/content/sample_data/'+fn, '/content/image')\n",
        "    break\n",
        "  %cd /content"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uDaxVZk04PI5",
        "cellView": "form"
      },
      "source": [
        "#@title Optionally shorten Driver video\n",
        "start_seconds =  0#@param {type:\"number\"}\n",
        "duration_seconds =  10#@param {type:\"number\"}\n",
        "start_seconds = max(start_seconds,0)\n",
        "duration_seconds = max(duration_seconds,0)\n",
        " \n",
        "if duration_seconds:\n",
        "  !mv /content/video /content/full_video\n",
        "  !ffmpeg  -ss $start_seconds -t $duration_seconds -i /content/full_video -f mp4 /content/video -y"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Oxi6-riLOgnm",
        "cellView": "form"
      },
      "source": [
        "#@title Prepare assets\n",
        "#@markdown If you ran out of RAM this means that the video is too large. You can shorten it above.\n",
        " \n",
        "center_video_to_body = True #@param {type:\"boolean\"}\n",
        "crop_video_to_body = False #@param {type:\"boolean\"}\n",
        "video_crop_expansion_factor = 1.05 #@param {type:\"number\"}\n",
        "center_image_to_body = True #@param {type:\"boolean\"}\n",
        "crop_image_to_body = False #@param {type:\"boolean\"}\n",
        "image_crop_expansion_factor = 1.05 #@param {type:\"number\"}\n",
        "video_crop_expansion_factor = max(video_crop_expansion_factor, 1)\n",
        "image_crop_expansion_factor = max(image_crop_expansion_factor, 1)\n",
        " \n",
        "import imageio\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.animation as animation\n",
        "from skimage.transform import resize\n",
        "from IPython.display import HTML, clear_output\n",
        "import cv2\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        " \n",
        "hog = cv2.HOGDescriptor()\n",
        "hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())\n",
        " \n",
        "def fix_dims(im):\n",
        "    if im.ndim == 2:\n",
        "        im = np.tile(im[..., None], [1, 1, 3])\n",
        "    return im[...,:3]\n",
        " \n",
        "def get_crop(im, center_body=True, crop_body=True, expansion_factor=1, rects=None):\n",
        "    im = fix_dims(im)\n",
        "    if (center_body or crop_body) and rects is None:\n",
        "        rects, _ = hog.detectMultiScale(im, winStride=(4, 4),padding=(8,8), scale=expansion_factor)\n",
        "    if (center_body or crop_body) and rects is not None and len(rects):\n",
        "        x0,y0,w,h = sorted(rects, key=lambda x: x[2]*x[3])[-1]\n",
        "        if crop_body:\n",
        "            x0 += w//2-h//2\n",
        "            x1 = x0+h\n",
        "            y1 = y0+h\n",
        "        else:\n",
        "            img_h,img_w = im.shape[:2]\n",
        "            x0 += (w-img_h)//2\n",
        "            x1 = x0+img_h\n",
        "            y0 = 0\n",
        "            y1 = img_h\n",
        "    else:\n",
        "        h,w = im.shape[:2]\n",
        "        x0 = (w-h)//2\n",
        "        x1 = (w+h)//2\n",
        "        y0 = 0\n",
        "        y1 = h\n",
        "    return int(x0),int(x1),int(y0),int(y1)\n",
        " \n",
        "def pad_crop_resize(im, x0=None, x1=None, y0=None, y1=None, new_h=256, new_w=256):\n",
        "    im = fix_dims(im)\n",
        "    h,w = im.shape[:2]\n",
        "    if x0 is None:\n",
        "      x0 = 0\n",
        "    if x1 is None:\n",
        "      x1 = w\n",
        "    if y0 is None:\n",
        "      y0 = 0\n",
        "    if y1 is None:\n",
        "      y1 = h\n",
        "    if x0<0 or x1>w or y0<0 or y1>h:\n",
        "        im = np.pad(im, pad_width=[(max(-y0,0),max(y1-h,0)),(max(-x0,0),max(x1-w,0)),(0,0)], mode='edge')\n",
        "    return resize(im[max(y0,0):y1-min(y0,0),max(x0,0):x1-min(x0,0)], (new_h, new_w))\n",
        " \n",
        " \n",
        "source_image = imageio.imread('/content/image')\n",
        "source_image = pad_crop_resize(source_image, *get_crop(source_image, center_body=center_image_to_body, crop_body=crop_image_to_body, expansion_factor=image_crop_expansion_factor), new_h=res, new_w=res)\n",
        " \n",
        "with imageio.get_reader('/content/video', format='mp4') as reader:\n",
        "  fps = reader.get_meta_data()['fps']\n",
        " \n",
        "  driving_video = []\n",
        "  rects = None\n",
        "  try:\n",
        "      for i,im in enumerate(reader):\n",
        "          if not crop_video_to_body:\n",
        "              break\n",
        "          rects, _ = hog.detectMultiScale(im, winStride=(4, 4),padding=(8,8), scale=video_crop_expansion_factor)\n",
        "          if rects is not None and len(rects):\n",
        "              break\n",
        "      x0,x1,y0,y1 = get_crop(im, center_body=center_video_to_body, crop_body=crop_video_to_body, expansion_factor=video_crop_expansion_factor, rects=rects)\n",
        "      reader.set_image_index(0)\n",
        "      for im in reader:\n",
        "          driving_video.append(pad_crop_resize(im,x0,x1,y0,y1,new_h=res,new_w=res))\n",
        "  except RuntimeError:\n",
        "      pass\n",
        " \n",
        "def vid_display(source, driving, generated=None):\n",
        "    fig = plt.figure(figsize=(8 + 4 * (generated is not None), 6))\n",
        " \n",
        "    ims = []\n",
        "    for i in range(len(driving)):\n",
        "        cols = [source]\n",
        "        cols.append(driving[i])\n",
        "        if generated is not None:\n",
        "            cols.append(generated[i])\n",
        "        im = plt.imshow(np.concatenate(cols, axis=1), animated=True)\n",
        "        plt.axis('off')\n",
        "        ims.append([im])\n",
        " \n",
        "    ani = animation.ArtistAnimation(fig, ims, interval=50, repeat_delay=1000)\n",
        "    plt.close()\n",
        "    return ani\n",
        " \n",
        "clear_output()\n",
        "if rects is not None and len(rects):\n",
        "    print('first found body in frame %d'%i)\n",
        "HTML(vid_display(source_image, driving_video).to_html5_video())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SB12II11kF4c",
        "cellView": "form"
      },
      "source": [
        "#@title Animate\n",
        " \n",
        "animation_mode = 'avd' #@param ['standard', 'relative', 'avd']\n",
        "\n",
        "%cd /content/articulated-animation\n",
        "import torch\n",
        "from tqdm import tqdm\n",
        "from skimage import img_as_ubyte\n",
        "from animate import get_animation_region_params\n",
        " \n",
        "def make_animation(source_image, driving_video, generator, region_predictor, avd_network,\n",
        "                   animation_mode='standard', cpu=False):\n",
        "    with torch.no_grad():\n",
        "        predictions = []\n",
        "        source = torch.tensor(source_image[np.newaxis].astype(np.float32)).permute(0, 3, 1, 2)\n",
        "        if not cpu:\n",
        "            source = source.cuda()\n",
        "        driving = torch.tensor(np.array(driving_video)[np.newaxis].astype(np.float32)).permute(0, 4, 1, 2, 3)\n",
        "        source_region_params = region_predictor(source)\n",
        "        driving_region_params_initial = region_predictor(driving[:, :, 0])\n",
        "\n",
        "        for frame_idx in tqdm(range(driving.shape[2])):\n",
        "            driving_frame = driving[:, :, frame_idx]\n",
        "            if not cpu:\n",
        "                driving_frame = driving_frame.cuda()\n",
        "            driving_region_params = region_predictor(driving_frame)\n",
        "            new_region_params = get_animation_region_params(source_region_params, driving_region_params,\n",
        "                                                            driving_region_params_initial, avd_network=avd_network,\n",
        "                                                            mode=animation_mode)\n",
        "            out = generator(source, source_region_params=source_region_params, driving_region_params=new_region_params)\n",
        "\n",
        "            predictions.append(np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0])\n",
        "    return predictions\n",
        "\n",
        "predictions = make_animation(source_image, driving_video, generator, region_predictor, avd_network,\n",
        "                              animation_mode=animation_mode) \n",
        "imageio.mimsave('/content/generated.mp4', [img_as_ubyte(frame) for frame in predictions], fps=fps)\n",
        "!ffmpeg -i /content/generated.mp4 -i /content/video -c:v libx264 -c:a aac -map 0:v -map 1:a? -pix_fmt yuv420p -profile:v baseline -movflags +faststart /content/final.mp4 -y\n",
        " \n",
        "clear_output()\n",
        "HTML(vid_display(source_image, driving_video, predictions).to_html5_video())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lH4oJcBOKZcf",
        "cellView": "form"
      },
      "source": [
        "#@title Download\n",
        "\n",
        "from google.colab import files\n",
        "files.download('/content/final.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}
